from src.ga.interpreter import str2tensor
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable


def get_torch_output(tensor, operator, variable=None):
    real_tensor = torch.Tensor(str2tensor(tensor))
    if operator == 'avg_pool':
        return F.avg_pool2d(real_tensor, kernel_size=2, stride=2)
    elif operator == 'avg_pool3d':
        return F.avg_pool3d(real_tensor, 2, 2, 'SAME')
    elif operator == 'bias_add':
        bias_nd = np.broadcast_to(variable, real_tensor.shape)
        bias_nd.setflags(write=True)
        my_bias = torch.Tensor(bias_nd)
        return real_tensor.add(my_bias)
    elif operator == 'conv1d':
        # my_filter = tf.Variable(tf.random.normal([1, real_tensor.shape[-1], 1], stddev=0.01))
        conv1d = torch.nn.Conv1d(real_tensor.shape[-1], 1, 1)
        variable.setflags(write=True)
        param = torch.nn.Parameter(torch.Tensor(variable))
        conv1d.weight = param
        return conv1d(real_tensor)
    # elif operator == 'conv2d':
    #     # my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1], 32], stddev=0.01))
    #     conv1d = torch.nn.Conv1d(real_tensor.shape[-1], 1, 1)
    #     variable.setflags(write=True)
    #     param = torch.nn.Parameter(torch.Tensor(variable))
    #     conv1d.weight = param
    #     return conv1d(real_tensor)
#     elif operator == 'conv3d':
#         my_filter = tf.Variable(tf.random.normal([3, 3, 3, real_tensor.shape[-1], 32], stddev=0.01))
#         return tf.nn.conv3d(real_tensor, my_filter, strides=[1, 1, 1, 1, 1], padding='SAME')
#     elif operator == 'dilation2d':
#         my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1]], stddev=0.01))
#         return tf.nn.dilation2d(real_tensor, my_filter, strides=[1, 1, 1, 1], padding='SAME',
#                                 data_format='NHWC', dilations=[1, 1, 1, 1])
#     elif operator == 'depthwise_conv2d':
#         my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1], 32], stddev=0.01))
#         return tf.nn.depthwise_conv2d(real_tensor, my_filter, strides=[1, 1, 1, 1], padding='SAME')
#     elif operator == 'softmax':
#         return tf.nn.softmax(real_tensor)
#     elif operator == 'erosion2d':
#         my_filter = tf.Variable(tf.random.normal([3, 3, real_tensor.shape[-1]], stddev=0.01))
#         return tf.nn.depthwise_conv2d(real_tensor, my_filter, strides=[1, 1, 1, 1], padding='SAME',
#                                       data_format='NHWC', dilations=[1, 1, 1, 1])
#     elif operator == 'log_softmax':
#         return tf.nn.log_softmax(real_tensor)
#     else:
#         return ''
# w=torch.nn.Conv2d(2,2,3,padding=1)
# print(w.weight)
# print(type(w.weight))
# a = np.ones(w.weight.shape)
# aw = torch.nn.Parameter(torch.Tensor(a))
# w.weight = aw
# print(w.weight)